#define WIN32_LEAN_AND_MEAN
#include <windows.h>
#include <winsock2.h>
#include <ws2tcpip.h>
#include <stdio.h>
#include <string.h>
#include <stdlib.h>

// Disable security warnings
#pragma warning(disable: 4996)

int main(int argc, char* argv[])
{
    WSADATA wsaData;
    SOCKET sock;
    struct sockaddr_in addr;
    int addrLen = sizeof(addr);
    unsigned char buffer[2048] = { 0 };
    int port = 50000;
    char senderIP[INET_ADDRSTRLEN]; // Declare senderIP here

    // Parse command line arguments for port number
    if (argc > 1)
    {
        port = atoi(argv[1]);
        if (port <= 0 || port > 65535)
        {
            printf("Invalid port number: %s\n", argv[1]);
            printf("Usage: wollisten [port]\n");
            printf("Default port: 50000\n");
            return 1;
        }
    }

    // Initialize Winsock
    if (WSAStartup(MAKEWORD(2, 2), &wsaData) != 0)
    {
        printf("Winsock initialization failed\n");
        return 1;
    }

    // Create UDP socket
    sock = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
    if (sock == INVALID_SOCKET)
    {
        printf("Failed to create socket\n");
        WSACleanup();
        return 1;
    }

    // Set socket options to allow address reuse
    int optval = 1;
    if (setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, (char*)&optval, sizeof(optval)) != 0)
    {
        printf("Failed to set socket options\n");
        closesocket(sock);
        WSACleanup();
        return 1;
    }

    // Bind socket to port
    memset(&addr, 0, sizeof(addr));
    addr.sin_family = AF_INET;
    addr.sin_addr.s_addr = htonl(INADDR_ANY);
    addr.sin_port = htons(port);

    if (bind(sock, (struct sockaddr*)&addr, sizeof(addr)) == SOCKET_ERROR)
    {
        printf("Failed to bind to port: %d\n", port);
        closesocket(sock);
        WSACleanup();
        return 1;
    }

    printf("Listening for WOL packets...\n");
    printf("Listening port: %d\n", port);
    printf("Press Ctrl+C to exit\n\n");

    // Main loop to receive packets
    while (true)
    {
        int recvLen = recvfrom(sock, (char*)buffer, sizeof(buffer), 0, (struct sockaddr*)&addr, &addrLen);
        if (recvLen > 0)
        {
            // Check if this is a WOL packet (starts with 6 bytes of 0xFF)
            if (recvLen >= 102 && buffer[0] == 0xFF && buffer[1] == 0xFF && buffer[2] == 0xFF &&
                buffer[3] == 0xFF && buffer[4] == 0xFF && buffer[5] == 0xFF)
            {
                // Extract MAC address
                unsigned char mac[6];
                memcpy(mac, &buffer[6], 6);

                // Get sender's IP address and port
                inet_ntop(AF_INET, &(addr.sin_addr), senderIP, INET_ADDRSTRLEN);
                int senderPort = ntohs(addr.sin_port);

                // Display information about the received WOL packet
                printf("Received WOL packet!\n");
                printf("MAC address: %02X:%02X:%02X:%02X:%02X:%02X\n",
                    mac[0], mac[1], mac[2], mac[3], mac[4], mac[5]);
                printf("Sender: %s:%d\n", senderIP, senderPort);
                printf("Packet size: %d bytes\n\n", recvLen);
            }
        }
        else
        {
            // Receive failed, ignore error and continue
            Sleep(10);
        }
    }

    // Cleanup resources
    closesocket(sock);
    WSACleanup();

    return 0;
}